#include "timeIntegrator.h"

using namespace std;

double ForwardEuler::coefficients[3][1]    = {1.,1.,0.};
double RungeKutta22::coefficients[4][2]    = {{0.0,0.0},{1.0,0.0},{0.5,0.5},{0.0,1.0}};
double RungeKutta33::coefficients[5][3]    = {{0.0,0.0,0.0},{1./2.,0.0,0.0},{-1.0,2.0,0.0},{1./6.,2./3.,1./6.},{0.0,1./2.,1.}};   
double RungeKutta44::coefficients[6][4]    = {{0.0,0.0,0.0,0.0},{0.5,0.0,0.0,0.0},{0.0,0.5,0.0,0.0},{0.0,0.0,1.0,0.0},{1./6.,1./3.,1./3.,1./6.},{0.0,0.5,0.5,1.}};
double SSPRungeKutta22::coefficients[4][2] = {{0.0,0.0},{1.0,0.0},{0.5,0.5},{0.0,1.0}};
double SSPRungeKutta32::coefficients[5][3] = {{0.0,0.0,0.0},{1./2.,0.0,0.0},{1./2.,1./2.,0.0},{1./3.,1./3.,1./3.},{0.0,1./2.,1.}};
double SSPRungeKutta33::coefficients[5][3] = {{0.0,0.0,0.0},{1.,0.0,0.0},{1./4.,1./4.,0.0},{1./6.,1./6.,2./3.},{0.0,1.,1./2.}};

//------------------------------------------------------------------------------
TimeIntegrator::TimeIntegrator(Model* mymodel,ClassParameter& param){

  model      = mymodel;
  nbTimeStep = param.getInteger("NbStep");
  dt         = param.getDouble("TimeStep");
  outputFreq = param.getInteger("OutputFrequency");
  Step       = 1;
  realTime   = 0;
  
  cout << "[TScheme] - nb time step -------> " << nbTimeStep << endl;	
  cout << "[TScheme] - time step ----------> " << dt         << endl;	
  cout << "[Output]  - frequency ----------> " << outputFreq << endl;
  
}
//------------------------------------------------------------------------------
void TimeIntegrator::advanceInTime(InitialSolution* InitSol){
    
  SolContainer* Uinit = InitSol->getInitialSol();
  
  // --- define containers -----------------------------------------------------
  
  SolContainer          DGSolution;
  SolContainer          tmpSol;
  vector<SolContainer>  stepSol;
  DGMatrix              MPminmax(model->mesh->getNbElmt(0),3);                  // (eID,3) --> [eID,min,max,MPV]     //NOT YET IMPLEMENTED for quad
 
  int eT=0;
  for(SolContainer::const_iterator eTIt=Uinit->begin();eTIt!=Uinit->end();++eTIt,eT++){
    int nbNP   = (*model->fs)[eT]->getNbENP();
    int nbElmt = model->mesh->getNbElmt(eT);
    DGSolution.push_back(DGMatrix(nbNP,nbElmt));
    tmpSol.push_back(DGMatrix(nbNP,nbElmt));
  } 
  
  for(int i=0;i<nbSubStep();i++){
    SolContainer tmp;
    int eT=0;
    for(SolContainer::const_iterator eTIt=Uinit->begin();eTIt!=Uinit->end();++eTIt,eT++){
      int nbNP   = (*model->fs)[eT]->getNbENP();
      int nbElmt = model->mesh->getNbElmt(eT);
      tmp.push_back(DGMatrix(nbNP,nbElmt));
    }
    stepSol.push_back(tmp);
  }
  
  // --- copy initial solution -------------------------------------------------
  
  eT=0;
  for(SolContainer::const_iterator eTIt=Uinit->begin();eTIt!=Uinit->end();++eTIt,eT++){
    DGSolution[eT] = (*Uinit)[eT];
  } 
  
  model->computeMaxPrinciple(&DGSolution,&MPminmax);
  
  // ---------------------------------------------------------------------------

  double tstart = omp_get_wtime();

  // --- TimeStep Loop ---------------------------------------------------------
  
  for(int iStep=1;iStep<nbTimeStep+1;iStep++,Step++){
    
    // --- update time ---------------------------------------------------------
  
    realTime += dt;
    
    // --- SubStep Loop --------------------------------------------------------
    
    for(int iSubStep=0;iSubStep<nbSubStep();iSubStep++){
      
      int eT=0;
      for(SolContainer::const_iterator eTIt=DGSolution.begin();eTIt!=DGSolution.end();++eTIt){
        tmpSol[eT] = DGSolution[eT];
      }
 
      for(int j=0;j<iSubStep;j++){
        if(aij(iSubStep,j)!=0){
          int eT=0;
          for(SolContainer::const_iterator eTIt=DGSolution.begin();eTIt!=DGSolution.end();++eTIt){
            tmpSol[eT].scaleAndAdd(aij(iSubStep,j)*dt,stepSol[j][eT]);
          }
        }
      }
      
      model->computeResidual(&tmpSol,&stepSol[iSubStep]);

    }
    
    // --- Update solution -----------------------------------------------------
    
    for(int iSubStep=0;iSubStep<nbSubStep();iSubStep++){
      int eT=0;
      for(SolContainer::const_iterator eTIt=DGSolution.begin();eTIt!=DGSolution.end();++eTIt){
        DGSolution[eT].scaleAndAdd(bi(iSubStep)*dt,stepSol[iSubStep][eT]);
      }
    }
    
    model->computeMaxPrinciple(&DGSolution,&MPminmax);
    
    // --- Save solution -------------------------------------------------------
    
    if(iStep%outputFreq==0){
      model->output->saveOutput(&DGSolution,iStep,realTime,1,"solution");
      model->output->saveOutput(&MPminmax,iStep,realTime,1,"MaximumPrincipleViolation"); 
    }    
    // --- Print infos ---------------------------------------------------------
    
    printf("[Solver]: Step = %i | CFLmax = %1.4lf | Time(s) = %g \n",iStep,model->getCFL(dt),realTime);
    
  }
   
  DGVector* Lnorm = computeLnorm(&DGSolution,InitSol);
  printf("[Solver]: L1 error   = %1.5e\n",(*Lnorm)(0));
  printf("[Solver]: L2 error   = %1.5e\n",(*Lnorm)(1));
  printf("[Solver]: LInf error = %1.5e\n",(*Lnorm)(2));
  
  double tend = omp_get_wtime();
  printf ("[Solver]: elapsed time is %.6lf seconds\n",tend-tstart);
  
}
//------------------------------------------------------------------------------
DGVector* TimeIntegrator::computeLnorm(const SolContainer* U,InitialSolution* InitSol){
  
  model->computeUQP(U);
  const SolContainer* UEQP = model->getUEQP();
  
  DGVector* Lnorm = new DGVector(3);                                            // (L1norm,L2norm)
  
  int eT=0;
  for(FsContainer::const_iterator fsIt=model->fs->begin();fsIt!=model->fs->end();++fsIt,eT++){
    
    int nbElmt                    = model->mesh->getNbElmt(eT);
    int nbEQP                     = (*fsIt)->getNbEQP();
    int nbF                       = model->mesh->getNbElmtFaces(eT);
    const DGMatrix* EQP           = (*fsIt)->getEQP();
    const DGVector* w             = (*fsIt)->getEW();
    FunctionSpace* mapping        = FunctionSpaceFactory::CreateNewFunctionSpace((FctType)model->mesh->getEGmshType(eT),1);
    
    //#pragma omp parallel for
    for(int eID=0;eID<nbElmt;eID++){
      double detJ = model->mesh->getDetJ(eT,eID);
      DGMatrix EQPcoord(nbEQP,2);
      for(int q=0;q<nbEQP;q++){
        for(int j=0;j<nbF;j++){
          double phi_j   = mapping->getPhiE(j,(*EQP)(q,0),(*EQP)(q,1)); 
          EQPcoord(q,0) += model->mesh->getNode(model->mesh->getElmtNodesTag(eT,eID,j),0)*phi_j;
          EQPcoord(q,1) += model->mesh->getNode(model->mesh->getElmtNodesTag(eT,eID,j),1)*phi_j;
        }
        double u        = (*UEQP)[eT](q,eID);
        double analytic = InitSol->analyticalSol(EQPcoord(q,0),EQPcoord(q,1));
        (*Lnorm)(0) += (*w)(q)*detJ*fabs(u-analytic);
        (*Lnorm)(1) += (*w)(q)*detJ*(u-analytic)*(u-analytic);
        (*Lnorm)(2)  = max((*Lnorm)(2),fabs(u-analytic));
      }
    }
    
    delete mapping;
    
  }

  (*Lnorm)(1) = sqrt((*Lnorm)(1));
  
  return Lnorm;
  
}
//------------------------------------------------------------------------------
TimeIntegrator* TimeIntegratorFactory::CreateNewTimeIntegrator(Model* mymodel,ClassParameter& param) {
      
  string type = param.getToken("TimeScheme");
  cout << "[TScheme] - type ---------------> " << type << endl;	
  
  if(type=="FE")    return new ForwardEuler(param,mymodel);
  if(type=="RK22")  return new RungeKutta22(param,mymodel);
  if(type=="RK33")  return new RungeKutta33(param,mymodel);
  if(type=="RK44")  return new RungeKutta44(param,mymodel);
  if(type=="SSP22") return new RungeKutta44(param,mymodel);
  if(type=="SSP32") return new RungeKutta44(param,mymodel);
  if(type=="SSP33") return new RungeKutta44(param,mymodel);
         
}

